from milvus_service import VectorDatabase
from ai_service import SiliconflowService
from dotenv import load_dotenv
import os
from pymilvus import DataType, FieldSchema, Collection, utility


def connect_to_vector_db(collection_name="product", host="localhost", port="19530"):
    """连接已有商品推荐向量数据库

    Args:
        collection_name: 集合名称
        host: Milvus服务地址
        port: Milvus服务端口
    Returns:
        VectorDatabase实例
    """
    # 初始化向量数据库
    vector_db = VectorDatabase(host=host, port=port)
    vector_db.connect()

    # 检查集合是否存在
    if collection_name in utility.list_collections():
        vector_db.collection = Collection(collection_name)
        print(f"成功连接到集合: {collection_name}")
    else:
        print(f"集合 {collection_name} 不存在，请先创建数据库")

    return vector_db


def create_vector_db(collection_name="product", host="localhost", port="19530"):
    """创建商品推荐向量数据库

    Args:
        collection_name: 集合名称
        host: Milvus服务地址
        port: Milvus服务端口
    Returns:
        VectorDatabase实例
    """
    # 定义集合字段结构
    fields = [
        FieldSchema(
            name="title",
            dtype=DataType.VARCHAR,
            is_primary=True,
            max_length=100,
        ),
        FieldSchema(name="description", dtype=DataType.VARCHAR, max_length=500),
        FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=384),
    ]

    # 初始化向量数据库
    vector_db = VectorDatabase(host=host, port=port)
    vector_db.connect()
    vector_db.drop_collection(collection_name)
    vector_db.create_collection(collection_name, fields, "商品推荐知识库")

    # 插入数据
    titles = [
        "智能手机X20",
        "无线降噪耳机Max",
        "4K超高清智能电视",
        "迷你便携式投影仪",
        "智能健康手表S7",
        "机械背光键盘G90",
        "光学游戏鼠标M5",
        "2K屏平板电脑P8",
        "彩色家用打印机P30",
        "虚拟环绕声电竞耳机H9",
    ]
    descriptions = [
        "搭载顶级处理器，具备超强性能和流畅体验的智能手机，支持5G网络和快速充电功能",
        "采用先进降噪技术，提供沉浸式音乐体验和清晰通话质量的无线蓝牙耳机",
        "拥有4K超高清分辨率和智能交互系统，带来震撼视觉效果的智能电视",
        "体积小巧轻便，便于携带和随时随地投影使用的迷你投影仪",
        "集成多种健康监测功能，如心率、睡眠监测等的多功能智能手表",
        "配备背光设计，按键手感舒适且反馈灵敏的机械键盘，适合游戏和办公使用",
        "采用高精度光学传感器，实现精准操控和快速响应的游戏鼠标",
        "配备2K高清屏幕，轻薄便携，适合办公、娱乐和学习的平板电脑",
        "支持彩色打印，具备高效打印速度和出色打印质量的家用打印机",
        "采用虚拟环绕声技术，提供沉浸式游戏音效体验的电竞耳机",
    ]
    vector_db.insert_data(
        vector_mapping={"vector": "description"},
        title=titles,
        description=descriptions,
    )

    # 创建索引
    vector_db.create_index("vector")

    return vector_db


def product_ai_chat(question, vector_db, ai_service):
    """商品推荐智能问答流程

    Args:
        question: 用户问题
        vector_db: VectorDatabase实例
        ai_service: AI服务实例
    Returns:
        str: AI生成的回答
    """
    # 1. 向量数据库检索
    search_results = vector_db.search(
        query_text=question,
        vector_field="vector",
        output_fields=["title", "description"],
        limit=3,
    )

    # 2. 如果没有搜索结果，直接返回不知道
    if not search_results:
        print("-" * 50)
        content = "抱歉，我没有找到相关商品信息。"
        print(content)
        return content

    # 3. 构建AI提示
    context = "\n".join(
        [
            f"[搜索结果] 第{search_result['rank']}名: {search_result['title']}，向量距离：{search_result['distance']}，商品描述：{search_result['description']}，相似度: {search_result['similarity']:.2f}"
            for search_result in search_results
        ]
    )
    print(f"搜索结果: \n{context}")

    prompt = f"""根据以下商品搜索结果回答用户问题:
    {context}

    用户问题: {question}
    请用中文回答，保持专业和友好，不要提到搜索相关的文字。如果不知道，请说“抱歉，我没有找到相关商品信息。”"""

    # 4. 调用AI服务生成回答
    print("-" * 50)
    print("AI回复内容:")
    response = ai_service.ai_chat(prompt)
    return response


if __name__ == "__main__":
    # 加载环境变量
    load_dotenv()

    # 获取API密钥
    api_key = os.getenv("SILICONFLOW_API_KEY")

    # 初始化AI服务
    ai_service = SiliconflowService(api_key)

    # 选择连接已有数据库或创建新数据库
    mode = input("请选择模式(1-创建新数据库, 2-连接已有数据库): ")
    if mode == "1":
        vector_db = create_vector_db()
    else:
        vector_db = connect_to_vector_db()

    # 测试商品推荐智能问答流程
    question = "我想要一个降噪的耳机，有什么推荐吗"
    answer = product_ai_chat(question, vector_db, ai_service)
